import scanpy as sc
import pandas as pd
import numpy as np
import json
import os
import shutil


def convert_visium_hd_to_h5ad(bin_dir, output_path="visium_hd_dense.h5ad"):
    """
    将 Visium HD binned 数据转换为兼容的 h5ad 文件

    参数:
    bin_dir -- bin 数据目录路径 (e.g., "/root/autodl-tmp/binned_outputs/square_016um")
    output_path -- 输出的 h5ad 文件路径
    """
    # 1. 确保坐标文件格式正确
    spatial_dir = os.path.join(bin_dir, "spatial")
    parquet_path = os.path.join(spatial_dir, "tissue_positions.parquet")
    csv_path = os.path.join(spatial_dir, "tissue_positions_list.csv")

    if not os.path.exists(csv_path):
        position_df = pd.read_parquet(parquet_path)
        position_df.to_csv(csv_path, index=False, header=None)
        print(f"已转换坐标文件: {csv_path}")

    # 2. 使用 scanpy 读取 Visium 格式数据
    try:
        adata = sc.read_visium(bin_dir)
    except Exception as e:
        print(f"sc.read_visium 错误: {e}")
        # 手动读取作为备选方案
        adata = sc.read_10x_mtx(
            os.path.join(bin_dir, "filtered_feature_bc_matrix"),
            var_names='gene_symbols',
            make_unique=True
        )

        # 手动添加空间坐标
        positions = pd.read_csv(csv_path, header=None, index_col=0)
        positions.columns = [
            'in_tissue', 'array_row', 'array_col',
            'pxl_col_in_fullres', 'pxl_row_in_fullres'
        ]
        positions = positions.loc[adata.obs_names]
        adata.obsm['spatial'] = positions[['pxl_row_in_fullres', 'pxl_col_in_fullres']].values

    # 3. 添加元数据
    adata.var_names_make_unique()

    # 4. 保存结果
    adata.write(output_path)
    print(f"成功保存 h5ad 文件到: {output_path}")
    print(f"数据维度: {adata.shape[0]} 个点, {adata.shape[1]} 个基因")

    return adata


# 使用示例
if __name__ == "__main__":
    bin_dir = "square_016um"
    output_h5ad = "/visium_hd_dense.h5ad"

    adata = convert_visium_hd_to_h5ad(
        bin_dir=bin_dir,
        output_path=output_h5ad
    )